#!/usr/bin/env python3
"""
为 Markdown 协议表格批量新增“功能码”列，并根据每个表格上方最近的小标题自动补充。
用法：python3 scripts/add_function_code_column.py config/protocols/monet/monet.md
"""
import sys, re

if len(sys.argv) < 2:
    print("usage: add_function_code_column.py file.md")
    sys.exit(2)

path = sys.argv[1]
with open(path, 'r', encoding='utf-8') as f:
    lines = f.readlines()

out = []
changed = False
current_func = ''
# 匹配如“功能码 03/04”或“功能码03/04”
func_pat = re.compile(r'功能码\s*([0-9/]+)')
# 匹配表头
header_pat = re.compile(r'^\s*\|\s*地址\s*\|')
# 匹配分隔行
sep_pat = re.compile(r'^\s*\|[\s\-:|]+\|\s*$')
# 匹配数据行
row_pat = re.compile(r'^\s*\|.*\|\s*$')

i = 0
while i < len(lines):
    line = lines[i]
    # 检查小节标题中的功能码
    m = func_pat.search(line)
    if m:
        current_func = m.group(1)
    # 检查表头
    if header_pat.match(line):
        # 插入功能码列
        cols = [c.strip() for c in line.strip().split('|') if c.strip()]
        if '功能码' not in cols:
            idx = cols.index('备注') if '备注' in cols else len(cols)
            new_cols = []
            for c in cols:
                new_cols.append(c)
                if c == 'bit':
                    new_cols.append('功能码')
            # 兼容没有 bit 列的表格
            if 'bit' not in cols:
                new_cols.insert(-1, '功能码')
            new_header = '| ' + ' | '.join(new_cols) + ' |\n'
            out.append(new_header)
            changed = True
        else:
            out.append(line)
        # 处理分隔行
        i += 1
        sep = lines[i]
        ncol = out[-1].count('|')-1
        new_sep = '|' + '|'.join(['---']*ncol) + '|\n'
        out.append(new_sep)
        changed = True
        i += 1
        # 处理数据行
        while i < len(lines) and row_pat.match(lines[i]) and not header_pat.match(lines[i]):
            row = lines[i].strip()
            cells = [c.strip() for c in row.split('|') if c.strip()]
            # 跳过空行
            if not cells:
                out.append(lines[i])
                i += 1
                continue
            # 插入功能码列
            if 'bit' in cols:
                idx = cols.index('bit')+1
            else:
                idx = len(cells)-1
            # 若已有功能码列则跳过
            if len(cells) == ncol:
                out.append(lines[i])
                i += 1
                continue
            cells.insert(idx, current_func if current_func else '-')
            new_row = '| ' + ' | '.join(cells) + ' |\n'
            out.append(new_row)
            changed = True
            i += 1
        continue
    out.append(line)
    i += 1

if changed:
    bak = path + '.bak'
    with open(bak, 'w', encoding='utf-8') as f:
        f.writelines(lines)
    with open(path, 'w', encoding='utf-8') as f:
        f.writelines(out)
    print(f"Added 功能码 column to {path}; backup at {bak}")
else:
    print(f"No changes needed for {path}")
